package com.zx.lc.controller;

import com.zx.lc.config.SseUtil;
import com.zx.lc.mcp.ToolServicce;
import com.zx.lc.service.OllamaAssistant;
import com.zx.lc.service.OllamaAssistant2;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.service.TokenStream;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping("/chat")
public class ChatController {

    @Resource(name = "ollamaChatModel")
    private ChatLanguageModel chatLanguageModel;

    @Autowired
    private OllamaAssistant ollamaAssistant;

    @Autowired
    private OllamaAssistant2 ollamaAssistant2;

    @Autowired
    private ToolServicce toolServicce;

    @GetMapping("/chat")
    public String chat(@RequestParam(value = "message", defaultValue = "你是谁？") String message) {
        return chatLanguageModel.chat(message);
    }

    @GetMapping("/chat2")
    public String chat2(@RequestParam(value = "message", defaultValue = "你是谁？") String message) {
        return ollamaAssistant.chat(message);
    }

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter stream(@RequestParam(value = "message", defaultValue = "你是谁？") String message) {
        // 获得响应流
        TokenStream stream = ollamaAssistant.stream(message);

        return SseUtil.createSseEmitter(stream);
    }

    @GetMapping(value = "/stream2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter stream2(@RequestParam(value = "message", defaultValue = "你是谁？") String message,
            int memoryId) {
        // 获得响应流
        TokenStream stream = ollamaAssistant2.stream(memoryId, message);

        return SseUtil.createSseEmitter(stream);
    }


    @GetMapping(value = "/stream3", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> stream3(@RequestParam(value = "message", defaultValue = "你是谁？") String message) {
        // 获得响应流
        Flux<String> response = ollamaAssistant.stream3(message);
        return response;
    }

    @GetMapping("/tool1")
    public String tool1() {
        return toolServicce.chat("帮我计算 2+3 等于几");
    }

    @GetMapping("/tool2")
    public String tool2() {
        return toolServicce.chat("帮我获得zhangxue的网站地址");
    }

}